import json
import os

import pika
from encryption.EncryptionUtils import EncryptionUtils
from work.AssetsDetect import AssetsDetect
from work.threatDiscovery import ThreatDiscovery
from work.LogTransfer import LogTransfer

class RabbitMQ:
    def __init__(self):
        self.__host = '192.168.138.138'
        self.__port = '4568'
        self.__user = 'admin'
        self.__password = '123456'
        self.__virtual_host = 'my_vhost'
        self.__connection = ''
        self.__channel = ''
        # 初始化连接
        self.__get_connection()
        self.__encryptor = EncryptionUtils()
        # 初始化 log_transfer 属性
        self.log_transfer = None
        # 本地状态文件路径
        # 本地状态文件路径
        self.status_file = 'log_sync_status.json'
        # 读取状态文件
        self.log_sync_status = self.read_log_sync_status()

    def read_log_sync_status(self):
        """
        读取本地状态文件，获取日志同步状态
        :return: 日志同步状态（True 或 False）
        """
        if os.path.exists(self.status_file):
            try:
                with open(self.status_file, 'r') as f:
                    data = json.load(f)
                    print(f"成功读取状态文件: {data}")
                    return data.get('log_sync', False)
            except Exception as e:
                print(f"读取状态文件出错: {e}")
        return False

    def write_log_sync_status(self, status):
        """
        将日志同步状态写入本地文件
        :param status: 日志同步状态（True 或 False）
        """
        try:
            with open(self.status_file, 'w') as f:
                json.dump({'log_sync': status}, f)
        except Exception as e:
            print(f"写入状态文件出错: {e}")

    def __get_connection(self):
        """
        获取连接对象
        :return:
        """
        # 获取认证对象
        credentials = pika.PlainCredentials(self.__user, self.__password)
        # 获取连接
        self.__connection = pika.BlockingConnection(
            pika.ConnectionParameters(host=self.__host, port=self.__port, virtual_host=self.__virtual_host,
                                      credentials=credentials))
        # 获取通道
        self.__channel = self.__connection.channel()

    def __my_producer(self, exchange, routing_key, data):
        """
        生产者
        :param routing_key: 路由键
        :param exchange: 交换机
        :param data: 数据
        :return:
        """

        # 包加密
        data = self.__encryptor.aes_encrypt(data)
        #验证是否加密
        # print("数据"+data)
        self.__channel.basic_publish(exchange=exchange, routing_key=routing_key, body=data)

    def produce_sysinfo(self, data):
        """
        生产系统信息
        :return:
        """
        exchange = 'sysinfo_exchange'
        routing_key = 'sysinfo'
        # 发送数据
        self.__my_producer(exchange, routing_key, data)

    def produce_status_info(self, data):
        exchange = 'sysinfo_exchange'
        routing_key = 'status'
        #发送数据
        self.__my_producer(exchange, routing_key, data)

    def consume_queue(self, queue_name):
        """
        消费队列数据
        :param queue_name:
        :return:
        """

        # 消费队列
        self.__channel.basic_consume(queue=queue_name,on_message_callback=self.__process_message, auto_ack=True)
        # 开始监听
        self.__channel.start_consuming()

    def __process_message(self, ch, method, properties, message):
        """
        处理消息
        :param ch:
        :param properties:
        :param message:
        :return:
        """
        try:
            data = str(self.__encryptor.aes_decrypt(message))
            # print(data)
            data = json.loads(data)
            print(f"收到消息: {data}")

            # 解密数据

            # print(f"解密后消息: {data}")

            if data['type'] == 'assets':
                assets_detect = AssetsDetect(self, data)
                assets_detect.start()
            if data['type'] == 'threat':
                threat_discovery = ThreatDiscovery(self, data)
                threat_discovery.start()
            if data['type'] == 'logSync':
                if data['status'] == 1:
                    # 启动日志同步服务（如果未启动）
                    if not self.log_transfer or not self.log_transfer.is_running():
                        self.log_transfer = LogTransfer(self, data['macAddress'])
                        self.log_transfer.start_sync()
                    self.write_log_sync_status(True)
                elif data['status'] == 0:
                    # 停止日志同步服务
                    if self.log_transfer and self.log_transfer.is_running():
                        self.log_transfer.stop_sync()
                        self.log_transfer = None
                    self.write_log_sync_status(False)
                else:
                    print("未知的logSync状态")
        except json.JSONDecodeError as e:
            print(f"JSON解析失败: {e}, 原始消息: {message}")
        except KeyError as e:
            print(f"缺少必要字段: {e}")
        except Exception as e:
            print(f"处理消息失败: {e}")
        finally:
            # 如果是手动确认模式，这里应加上 ch.basic_ack
            pass

    def produce_account_info(self, data):
        exchange = 'sysinfo_exchange'
        routing_key = 'account'
        # 发送数据
        self.__my_producer(exchange, routing_key, data)

    def produce_service_info(self, data):
        """生产账户信息"""
        exchange = 'sysinfo_exchange'
        routing_key = 'service'
        self.__my_producer(exchange, routing_key, data)

    def produce_process_info(self, data):
        """生产账户信息"""
        exchange = 'sysinfo_exchange'
        routing_key = 'process'
        self.__my_producer(exchange, routing_key, data)
    def produce_app_info(self, data):
        """生产账户信息"""
        exchange = 'sysinfo_exchange'
        routing_key = 'app'
        self.__my_producer(exchange, routing_key, data)

    def produce_hotfix_info(self, data):
        """生产补丁信息"""
        exchange = 'sysinfo_exchange'
        routing_key = 'hotfix'
        self.__my_producer(exchange, routing_key, data)

    def produce_weak_pwd_info(self, data):
        """生产弱密码信息队列"""
        exchange = 'sysinfo_exchange'
        routing_key = 'weak_pwd'
        self.__my_producer(exchange, routing_key, data)

    def produce_vul_info(self, data):
        """生产漏洞信息队列"""
        exchange = 'sysinfo_exchange'
        routing_key = 'vul'
        self.__my_producer(exchange, routing_key, data)

    def produce_app_vul_info(self, data):
        """生产漏洞信息队列"""
        exchange = 'sysinfo_exchange'
        routing_key = 'app_vul'
        self.__my_producer(exchange, routing_key, data)


    def produce_system_vul_info(self, data):
        """生产漏洞信息队列"""
        exchange = 'sysinfo_exchange'
        routing_key = 'system_vul'
        self.__my_producer(exchange, routing_key, data)

    def produce_log_data(self, data):
        """发送日志信息到特定队列"""
        exchange = 'sysinfo_exchange'
        routing_key = 'logs'
        self.__my_producer(exchange, routing_key, data)

